# -*- coding: utf-8 -*-
"""
@Time    : 2025/2/12 20:07 
@Author  : ZhangShenao 
@File    : 1.精确缓存.py 
@Desc    : 精确缓存
"""
import os
from time import time

import dotenv
from langchain.globals import set_llm_cache
from langchain_core.caches import InMemoryCache
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()

# 设置全局缓存
# 使用InMemoryCache组件,实现内存缓存
set_llm_cache(InMemoryCache())

# 构造Chain
prompt = ChatPromptTemplate.from_messages([
    ("user", "{input}"),
])
llm = ChatOpenAI(
    openai_api_base=os.getenv("OPENAI_API_BASE"),
    model="gpt-4o-mini",
    temperature=0.2,
)
chain = prompt | llm | StrOutputParser()

# 调用2次,对比耗时
start_time = time()
response = chain.invoke({"input": "给我讲个一句话笑话"})
end_time = time()
print(response)
print(f"第一次调用耗时: {end_time - start_time}秒")

start_time = time()
response = chain.invoke({"input": "给我讲个一句话笑话"})
end_time = time()
print(response)
print(f"第二次调用耗时: {end_time - start_time}秒")

# 为什么程序员总是喜欢在海边工作？因为他们喜欢“海”量数据！
# 第一次调用耗时: 2.3311691284179688秒
# 为什么程序员总是喜欢在海边工作？因为他们喜欢“海”量数据！
# 第二次调用耗时: 0.0033271312713623047秒
